// Copyright (c) 2020 Presto Labs Pte. Ltd.
// Author: leon

#pragma once
#include <map>
#include <memory>
#include <mutex>

#include <boost/asio/io_service.hpp>
#include <boost/asio/post.hpp>
#include <boost/asio/strand.hpp>

#include "coin2/base/time.h"
#include "coin2/exchange/base/executor/topic_distributor.h"
#include "coin2/exchange/base/market/market_exchange_api.h"

namespace coin2::exchange::base::feed {
namespace impl {

using coin2::exchange::base::executor::TopicRecordPacket;

class LatencyExtractor : public std::enable_shared_from_this<LatencyExtractor> {
 public:
  LatencyExtractor() {}
  virtual ~LatencyExtractor() = default;
  virtual void ExtractLatency(std::shared_ptr<TopicRecordPacket> packet) = 0;

  void SetIoc(boost::asio::io_context::strand* strand) {
    strand_ = strand;
  }

  void PostExtractLatency(std::shared_ptr<TopicRecordPacket> packet) {
    boost::asio::post(
        *strand_,
        std::bind(
            &LatencyExtractor::ExtractLatency,
            shared_from_this(),
            packet));
  }

  virtual bool Latency(double* avg_latency) {
    std::scoped_lock _(mutex_);
    
    // no packet received, random latency is for several no packet received connections cases
    if (!last_packet_) {
      *avg_latency = GetCurrentTimestamp() + ((std::rand() % 10) * 1'000'000'000LL);
      return true;
    }

    // receive packet without exchange time
    if (count_ == 0) {
      *avg_latency = GetCurrentTimestamp() - last_packet_->timestamp;
      return true;
    }

    // no update in this latency cycle
    if (last_latency_count_ == count_) {
      *avg_latency = GetCurrentTimestamp() - last_packet_->timestamp;
      return true;
    }

    *avg_latency = total_ft_et_diff_ / count_;
    last_latency_count_ = count_;
    return true;
  }

  std::shared_ptr<TopicRecordPacket> GetLastPacket() {
    std::scoped_lock _(mutex_);
    return last_packet_;
  }

 protected:
  void Update(double ft, double et) {
    if (et == 0) return; // no exchange time

    if (et < 0) {
      LOG_FIRST_N(WARNING, 10) << "Invalid et " << et;
      return;
    }

    if (ft < et) {
      LOG_FIRST_N(WARNING, 10) << "Invalid ft - et,  " << ft << " - " << et;
      return;
    }

    std::scoped_lock _(mutex_);
    // packet contains exchange time
    total_ft_et_diff_ += ft - et;
    ++count_;
  }

  void LogPacket(std::shared_ptr<TopicRecordPacket>& packet) {
    std::scoped_lock _(mutex_);
    last_packet_ = packet;
  }

  double total_ft_et_diff_ = 0;
  int64_t count_ = 0;
  int64_t last_latency_count_ = 0;
  // the data write from ws_connect io_context, read from ws_subscriber and health io_context
  std::mutex mutex_;
  boost::asio::io_context::strand* strand_;

  std::shared_ptr<TopicRecordPacket> last_packet_;
};

class TimeoutLatencyExtractor : public LatencyExtractor {
 public:
  TimeoutLatencyExtractor() {}
  virtual ~TimeoutLatencyExtractor() = default;
  virtual void ExtractLatency(std::shared_ptr<TopicRecordPacket> packet) override {
    LogPacket(packet);
  }

  virtual bool Latency(double* latency) override {
    std::scoped_lock _(mutex_);
    if (last_packet_) {
      *latency = GetCurrentTimestamp() - last_packet_->timestamp;
    } else {
      *latency = GetCurrentTimestamp();  // no data received 
    }
    return true;
  }
};

}  // namespace impl
using impl::LatencyExtractor;
using impl::TimeoutLatencyExtractor;
}  // namespace coin2::exchange::base::feed
